"""
特征提取
"""
import cv2
import numpy as np
from face_detect import FaceDetect
from net.facenet import InceptionResnetV1
import torch

device = torch.device('gpu:0' if torch.cuda.is_available() else 'cpu')


class FaceFeat:
    def __init__(self):
        # 初始化网络
        face_net = InceptionResnetV1().to(device)
        self.face_net = face_net
        # 加载参数
        face_net.load_state_dict(torch.load('weights/facenet_best_server.pt',map_location='cpu'))
        # 开启验证
        face_net.eval()
        pass

    def get_face_feat(self, face_img):
        """
        :param face_img: 人脸图像
        :return: 特征
        特征提取:
        1.因为图像(face_img)大小不一致，需要统一设置图像的大小h,w均为112
        2.归一化处理[-1, 1]
        3.因为Opencv打开的图像形状是HWC,而卷积神经网络输入数据形状是NCHW
           完成以下操作:
           换轴HWC-->CHW
           增加一个维度 (1, C,H,W)
        4.因为卷积神经网络输入的数据是张量tensor, 而当前你操作的数据还是numpy数组
        需要将Numpy数组转换为tensor -->(1, C, H, W) face_img_tensor
        5.使用facenet提取特征
        6. 怎么提取?
           face_feat_tensor = self.face_net(face_img_tensor)
        7. 因为后面很多的计算需要和numpy数组进行操作，需要转换
           face_feat_tensor--->face_feat
        """
        # 1.因为图像(face_img)大小不一致，需要统一设置图像的大小h,w均为112
        face_img = cv2.resize(face_img,dsize=(112,112))
        # 2.归一化处理[-1, 1]
        face_img = (face_img - 127.5)/127.5
        # 3.因为Opencv打开的图像形状是HWC,而卷积神经网络输入数据形状是NCHW
        #            完成以下操作:
        #            换轴HWC-->CHW
        #            增加一个维度 (1, C,H,W)
        face_img = np.transpose(face_img,(2,0,1))
        face_img = np.expand_dims(face_img,axis=0)
        # 需要将Numpy数组转换为tensor -->(1, C, H, W) face_img_tensor
        face_img_tensor = torch.Tensor(face_img).to(device)
        # 使用facenet提取特征
        face_feat_tensor = self.face_net(face_img_tensor)
        # 因为后面很多的计算需要和numpy数组进行操作，需要转换
        face_feat = face_feat_tensor.detach().cpu().numpy()

        return face_feat

if __name__ == '__main__':
    # 人脸检测
    frame = cv2.imread('images/test/sy_02.png')
    face_detect = FaceDetect()
    face_img = face_detect.detect(frame)
    # 特征提取
    face_feat = FaceFeat()
    feature = face_feat.get_face_feat(face_img)
    print(feature)